package com.cxyjy.cfyuiapi.websocket;

import com.alibaba.fastjson.JSONException;
import com.alibaba.fastjson.JSONObject;
import com.cxyjy.common.util.JObjectUtil;
import lombok.extern.slf4j.Slf4j;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;

import javax.websocket.Session;
import java.io.IOException;
import java.net.URI;
import java.util.HashMap;
import java.util.Map;

import static com.cxyjy.common.constants.CfyuiConstants.VEDIO_FACESWAP_DATA;

/**
 * @Author: jony
 * @Date : 2024/11/22 14:07
 * @VERSION v1.0
 */
@Slf4j
public class CfuiWebsocketClient extends WebSocketClient {
    private static Map<String, Object> countProgressMap = new HashMap<>();
    private int countValue;
    private int countMax;

    public CfuiWebsocketClient(URI serverUri) {
        super(serverUri);
    }

    @Override
    public void onOpen(ServerHandshake serverHandshake) {
        send("已连接到ComfyUI");
        countProgressMap = JObjectUtil.toMap(JObjectUtil.toMap(JSONObject.parse(VEDIO_FACESWAP_DATA)).get("prompt"));
        countMax = countProgressMap.size();
    }

    @Override
    public void onMessage(String msg) {
        //comfyUI的websocket返回的json有些是正常用双引号的，有些是只用斜杠的，要通过try catch来区分
        String type;
        Map<String, Object> msgMap = new HashMap<>();
        try {
            msgMap = JObjectUtil.toMap(JSONObject.parse(msg));
            type = JObjectUtil.stringValue(msgMap.get("type"));
        } catch (JSONException | ClassCastException e) {
            String msgReplace = msg.replace("\\", "\"");
            msgMap = JObjectUtil.toMap(JSONObject.parse(msgReplace));
            type = JObjectUtil.stringValue(msgMap.get("type"));
        }
        //硬件资源信息，后台用不上，不用管
        if (type.equals("crystools.monitor")) {
            return;
        }
        log.info(msg);
        //{"type": "executing", "data": {"node": null, "prompt_id": "5d80e6ba-4191-49f9-94ab-6b4cd7ddd1bc"}}
        if (type.equals("executing")) {
            Map<String, Object> dataMap = JObjectUtil.toMap(msgMap.get("data"));
            Object nodeObj = dataMap.get("node");
            String promptId = JObjectUtil.stringValue(dataMap.get("prompt_id"));
            Session session;
            WebSocket webSocket;
            try {
                session = WebSocket.getWebsocketSession(promptId);
                //发送关闭连接时的消息
                webSocket = (WebSocket) session.getUserProperties().get("websocketInstance");
                session.close();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
            //node为null时表示已经完成，关闭session
            if (JObjectUtil.isNull(nodeObj)) {
                closeConnection();
            } else {
                //用剩余节点数除以总结点数，取其补数传到前端
                String node = JObjectUtil.stringValue(nodeObj);
                countProgressMap.remove(node);
                countValue = countProgressMap.size();
                int progressPercent = 100 - countValue*100/countMax;
                Map<String,Object> resultMap = new HashMap<>();
                resultMap.put("progressPercent",progressPercent);
                webSocket.sendMessage(JSONObject.toJSONString(resultMap));
            }
        }
    }

    public void closeConnection() {
        close(1000, "换脸流程执行完毕，正常关闭连接");
    }


    @Override
    public void onClose(int code, String reason, boolean remote) {
        System.out.println("Connection closed: " + reason);
    }

    @Override
    public void onError(Exception e) {

    }
}
